#include "dnamap.h"
using namespace dnamap;

//---------------------------------------------------------------------
//class nucleotidebase
//---------------------------------------------------------------------

inline int nucleotidebase::codelookup(const char x)
{
    char t = toupper(x);
    switch(t)
    {
        case 'A': return 0;
        case 'T': return 1;
        case 'U': return 1;
        case 'G': return 2;
        case 'C': return 3;
    }
    return -1;
}

inline int nucleotidebase::dcodelookup(const char x)
{
    char t = toupper(x);
    switch(t)
    {
        case 'A': return 0x8;
        case 'T': return 0x4;
        case 'U': return 0x4;
        case 'G': return 0x2;
        case 'C': return 0x1;
        case 'K': return 0x6;
        case 'N': return 0xF;
        case 'M': return 0x9;
        case 'R': return 0xA;
        case 'Y': return 0x5;
        case 'S': return 0x3;
        case 'W': return 0xC;
        case 'B': return 0x7;
        case 'V': return 0xB;
        case 'H': return 0xD;
        case 'D': return 0xE;
        case 'X': return 0xF;
        case '.': return 0x0;
    }
    return 0;
}

inline int nucleotidebase::bcode2dcode(const int x)
{
    return 01 << (3 - x);
}

inline int nucleotidebase::dcodonlookup(const int x)
{
    int dcodon = 0, bcodon = x;
    for(int i = 0 ; i < codonlen; i++)
    {
        dcodon |= bcode2dcode(bcodon & 0x03) << (i * dbaselen);
        bcodon >>= 2;
    }
    return dcodon;
}

inline int nucleotidebase::dcodon2bcodon(const int x)
{
    int res = 0;
    for(int i = 0; i < codonlen; i++)
    {
        for(int j = 0; j < dbaselen; j++)
        {
            if(x & (01 << (j + i * dbaselen)))
            {
                res |= (dbaselen - j - 1) << (baselen * i);
                break;
            }
        }
    }
    return res;
}


//---------------------------------------------------------------------
//class dnaseq
//---------------------------------------------------------------------

//Default Constructor
dnaseq::dnaseq()
{
    seqlen = 0;
    framelen = 0;
}

dnaseq::~dnaseq()
{
}

//Constructor. Direct input DNA sequence as string X
dnaseq::dnaseq(const std::string x)
{
    seqlen = 0;
    framelen = 0;
    enter(x.c_str());
}

dnaseq::dnaseq(const std::string x, const int l)
{
    char writer[codonlen];
    framelen = x.length();
    seqlen = (framelen - 1) * codonlen;
    if(x[framelen - 1] != 0)
    {
        seqlen -= codonlen - x[framelen - 1];
    }
    readframe[0] = x;
    for(int i = 1; i < codonlen; i++)
    {
        writer[i] = 0;
        readframe[i].reserve(framelen);
    }
    for(int i = 0; i < framelen - 2; i++)
    {
        for(int j = 1; j < codonlen; j++)
        {
            char t = readframe[0][i] << (j * baselen);
            t |= readframe[0][i + 1] >> ((codonlen - j) * baselen);
            readframe[j].push_back(t);
        }
    }
    for(int i = 1; i < codonlen; i++)
    {
        readframe[i].push_back((seqlen - i) % codonlen);
    }
}

//Constructor. Direct input DNA sequence as char X
dnaseq::dnaseq(const char* x)
{
    seqlen = 0;
    framelen = 0;
    enter(x);
}

//Sequence input
void dnaseq::enter(const char* x)
{
    int i, dnalen = 0;         //Actual dna length
    char writer[codonlen];
    seqlen = strlen(x);
    framelen = seqlen / codonlen + 2;
    memset(writer, 0, sizeof(char) * codonlen);
    for(i = 0; i < codonlen; i++)
    {
        readframe[i].reserve(framelen);
    }

    for(i = 0; i < seqlen; i++)
    {
        int curcode = codelookup(x[i]);
        if(curcode >= 0)
        {
            dnalen ++;
            for(int j = 0; j < dnalen && j < codonlen; j++)
            {
                writer[j] |= curcode << ((codonlen - ((dnalen - j - 1) % codonlen) - 1) * baselen);
                //printf("%u %d %d %d\n", writer[0], dnalen - j, (dnalen - j) % codonlen, ((codonlen - ((dnalen - j) % codonlen)) * baselen));
                if((dnalen - j) % codonlen == 0)
                {
                    readframe[j].push_back(writer[j]);
                    //printf("%u %u\n", (unsigned int)(readframe[0][0]), writer[0]);
                    writer[j] = 0;
                }
            }

        }
    }
    for(i = 0; i < codonlen; i++)
    {
        if(((dnalen - i) % codonlen) != 0)
        {
            readframe[i].push_back(writer[i]);
        }
        readframe[i].push_back((dnalen - i) % codonlen);
    }
    seqlen = dnalen;
}

//Return base at position No.X, numeric
inline int dnaseq::baselookup(const int x)
{
    int pos = x / codonlen, off = codonlen - x % codonlen - 1;
    return (int(readframe[0].at(pos)) >> (off * baselen)) & basemask;
}

//Return base at position No. X, x is zero-based
char dnaseq::base(const int x)
{
    if(x < 0 || x > seqlen)
        return 0;
    return basecode[baselookup(x)];
}

//Return DNA sequence, char format
std::string dnaseq::nucleotide(const int start, const int len)
{
    int spos = start, epos = start + len;
    //Adjust start and length
    if(start < 0)
    {
        spos = 0;
    }
    if(epos > seqlen)
    {
        epos = seqlen;
    }

    std::string nucseq;
    nucseq.reserve(epos - spos + 2);                //Preserve space for higher efficiency
    for(int i = spos; i < epos; i++)
    {
        char curnt = basecode[baselookup(i)];
        nucseq.push_back(curnt);
    }
    nucseq += '\0';
    return nucseq;
}

//Reading frame retrieval
std::string dnaseq::getframe(const int num)
{
    return  readframe[num];
}

//Reading frame retrieval
std::string dnaseq::getframe(const int start, const int len)
{
    int s = start, num, sublen;
    if(s < 0)
    {
        s = 0;
    }
    if(s + len <= seqlen)
    {
        sublen = len;
    }
    else
    {
        sublen = seqlen - start;
    }

    std::string res = readframe[num].substr(start, sublen);
    res += char(sublen%codonlen);
    return res;
}

inline int dnaseq::length()
{
    return seqlen;
}

//---------------------------------------------------------------------
//class proteintranslator
//---------------------------------------------------------------------



protrans::protrans(const char* fname)
{
    FILE* fin = fopen(fname, "r");
    char charcodon[codonlen + 1];
    int amino;
    int ncodon;
    unsigned short codonrec[alphalen];
    memset(codonrec, 0, sizeof(short) * alphalen);
    memset(codoncount, 0 ,sizeof(int) * alphalen);

    while(feof(fin) == 0)
    {
        ncodon = amino = 0;
        fscanf(fin, "%s %c", charcodon, &amino);

        for(int i = 0; i < codonlen; i++)
        {
            int bcode = codelookup(charcodon[i]);
            if(bcode >= 0)
            {
                ncodon |= bcode << ((codonlen - i - 1) * 2);
            }
            else
            {
                ncodon = -1;
                break;
            }
        }

        if(ncodon >= 0)
        {
            //printf("%c %d\n", amino, amino);
            codon[ncodon] = amino;
            codoncount[amino]++;
            codonrec[amino] |= dcodonlookup(ncodon);
        }
        amino = fgetc(fin);
        //printf("%d\n", dcodonlookup(ncodon));
    }

    for(int i = 0; i < alphalen; i++)
    {
        if(codoncount[i] > 0)
        {
            anticodon[i] = new unsigned short[codoncount[i]];
            assert(anticodon[i]);
            codoncount[i] = 0;
        }
    }

    for(int i = 0; i < codonnum; i++)
    {
        DegenerateCodon[i] = codonrec[int(codon[i])];
        anticodon[int(codon[i])][codoncount[int(codon[i])]++] = dcodonlookup(i);
    }
}

std::string protrans::translate(std::string frame)
{
    std::string res = "";
    int len = frame.length();
    //printf("%d\n", frame[len - 1]);
    //printf("%d\n", frame[len - 2]);
    //printf("%d\n", frame[len - 3]);

    if(frame[len - 1] != 0)
    {
        len--;
    }
    len--;
    //printf("%d\n", len);
    res.reserve(len);
    for(int i = 0; i < len; i++)
    {
        if(frame[i] >= 0 && frame[i] < codonnum)
        {
            res += codon[int(frame[i])];
        }
    }
    return res;
}

sitelist protrans::synmut(std::string& frame, const char* dseq)
{
    int framelen = frame.length(), seqlen = strlen(dseq);           //reading frame length, replace site length
    int dframelen = 0;
    char* dframe = new char[framelen * codonlen];                   //denegerate frame
    sitelist sites(1);                                              //list of sites
    std::string tframe = frame;

    for(int i = 0; i < framelen - 1; i++)
    {
        unsigned short dcodon = DegenerateCodon[int(frame[i])];     //degenerate codon
        for(int j = 0; j < codonlen; j++)                           //save degenerate codon into single bases
        {
            dframe[dframelen + j] = (dcodon >> (dbaselen * (codonlen - j - 1))) & dbasemask;
        }
        dframelen += codonlen;
    }
    if(frame[framelen - 1] != 0)                                         //modify ending
    {
        dframelen += frame[framelen - 1] - codonlen;
        dframe[dframelen] = 0;
    }
    for(int i = 0; i <= dframelen - seqlen; i++)                    //Crude selection
    {
        int mismatch = 0;
        for(int j = 0; j < seqlen; j++)
        {
            if((dframe[i + j] & dseq[j]) == 0)
            {
                mismatch = 1;
                break;
            }
        }
        if(mismatch == 0)
        {
            sites.push_back(i);
        }
    }
    int i = 0;

    while(i < sites.length())
    {
        int codonseq = sites[i] / codonlen;                                         //Map to readframe
        char curaa = codon[frame[codonseq]];                                        //current amino acid
        int replacenum = 0;                                                         //replace sites
        int *replacepos = new int[seqlen / codonlen + 1];                           //position in frame
        unsigned short *replaceseq = new unsigned short[seqlen / codonlen + 1];     //new codon in numeric
        int replaceable = 0;                                                        //able to switch to fit
        for(int j = 0; j < seqlen; j++)
        {
            if(((bcode2dcode(tframe[sites[i] + j]) ^ dseq[j]) & (~dseq[j])) != 0)
            {
                replacepos[replacenum] = (sites[i] + j) / codonlen;
                for(int k = 0; k < codoncount[int(curaa)]; k++)
                {
                    int pos = sites[i] + j, rpos = j, match = 1;
                    rpos -= pos % codonlen;
                    pos = pos - (pos % codonlen);
                    for(int p = 0; p < codonlen; p++)
                    {
                        dframe[pos + p] = (anticodon[curaa][k] >>
                            ((codonlen - p - 1) * dbaselen)) & dbasemask;
                        if((rpos + p >= 0) &&
                            (((dframe[pos + p] ^ dseq[rpos + p]) & (~dseq[rpos + p])) != 0))
                        {
                            match = 0;
                            break;
                        }
                    }
                    if(match)
                    {
                        replaceseq[replacenum] = anticodon[curaa][k];
                        tframe.replace(replacepos[replacenum], 1, 1, dcodon2bcodon(anticodon[curaa][k]));
                        replacenum++;
                        replaceable = 1;
                        break;
                    }
                }
                if(replaceable == 0)
                {
                    break;
                }
            }
        }
        if(replaceable == 0)
        {
            sites.remove(i);
        }
        else
        {
            for(int j = 0; j < replacenum; j++)
            {
                frame.replace(replacepos[j], 1, 1, replaceseq[j]);
            }
            i++;
        }
    }
    return sites;
}

//---------------------------------------------------------------------
//class sitelist
//---------------------------------------------------------------------


sitelist::sitelist(const int x):stepspc(5)
{
    list = new int[x];
    assert(list);
    len = x;
    pos = -1;
}

sitelist::sitelist(const sitelist& x):stepspc(x.stepspc)
{
    len = x.len;
    pos = x.pos;
    list = new int[len];
    assert(list);
    memcpy(list, x.list, sizeof(int) * len);
}

sitelist&  sitelist::operator= (const sitelist& x)
{
    if(this != &x)
    {
        len = x.len;
        pos = x.pos;
        delete[] list;
        list = new int[len];
        assert(list);
        memcpy(list, x.list, sizeof(int) * len);
    }
    return *this;
}

sitelist::~sitelist()
{
    delete[] list;
}
int sitelist::length()
{
    return pos + 1;
}

int& sitelist::operator[] (const unsigned int x)
{
    assert(x >= 0 && x < len);
    return list[x];
}

void sitelist::remove(const int x)
{
    for(int i = x + 1; i < len; i++)
    {
        list[i - 1] = list[i];
    }
}

void sitelist::push_back(const int x)
{
    if(pos + 1 >= len)
    {
        int* nlist = new int[len + stepspc];
        assert(nlist);
        memcpy(nlist, list, sizeof(int) * len);
        nlist[pos + 1] = x;
        list = nlist;
        len += stepspc;
        pos ++;
    }
    else
    {
        list[++pos] = x;
    }
}

//---------------------------------------------------------------------
//class restriction
//---------------------------------------------------------------------

restriction::restriction(const char* rname, const char* rseq)
{
    sitename = new char[strlen(rname)];
    assert(sitename);
    strcpy(sitename, rname);
    sitelen = strlen(rseq);
    siteseq = new char[sitelen + 1];
    sitestr = new char[sitelen + 1];
    assert(siteseq);
    assert(sitestr);
    int pos = 0;
    for(int i = 0; i < sitelen; i++)
    {
        int dbasec = dcodelookup(rseq[i]);
        if(dbasec > 0)
        {
            siteseq[pos] =dbasec;
            sitestr[pos] = toupper(rseq[i]);
            pos ++;
        }
    }
    sitelen = pos;
    siteseq[pos] = sitestr[pos] = 0;
}

restriction::~restriction()
{
    delete[] sitename;
    delete[] siteseq;
}

